package be.fedict.eid.diagnostics.system.windows.jna;

import java.util.ArrayList;
import java.util.List;

import com.sun.jna.platform.win32.WinNT.HANDLE;
import com.sun.jna.ptr.IntByReference;
import com.sun.jna.ptr.PointerByReference;

import be.fedict.eid.diagnostics.system.data.Module;
import be.fedict.eid.diagnostics.system.data.SystemProcess;
import be.fedict.eid.diagnostics.system.windows.proxies.Kernel32;
import be.fedict.eid.diagnostics.system.windows.proxies.Psapi;
import be.fedict.eid.diagnostics.system.windows.proxies.Kernel32.LPPROCESSENTRY32;

public class ProcessHandler {
	
	private static final int TH32CS_SNAPPROCESS = (int) 0x00000002;
	private static final Kernel32 kernel32 = Kernel32.INSTANCE;
	private static final Psapi psapi = Psapi.INSTANCE;

	
	public static List<Module> fetchModules(SystemProcess proc) throws Exception {
		List<Module> modules = new ArrayList<Module>();
		
		if (!SecurityHandler.hasDebugPrivileges()) {
			if (!SecurityHandler.getDebugPrivileges());
				throw new Exception("Could not retrieve debug privileges. No modules available.");
		}
		
		// open process
		HANDLE hProcess = kernel32.OpenProcess(Kernel32.PROCESS_QUERY_INFORMATION | Kernel32.PROCESS_VM_READ, false, proc.getId());

		if (kernel32.GetLastError() == 5) // ERROR_ACCESS_DENIED
			throw new Exception("Could not open process '" + proc.getName() + "'.");
		
		// enumerate modules
		IntByReference lpcbNeededs = new IntByReference();
		HANDLE[] lphModule = new HANDLE[100*4];		
		if (psapi.EnumProcessModules(hProcess, lphModule, lphModule.length, lpcbNeededs)) {
			for (HANDLE handle : lphModule) {
				if (handle != null) {
					byte[] buffer = new byte[260];
					psapi.GetModuleBaseNameA(hProcess, handle, buffer, buffer.length);

					// create abstract representation
					Module mod = new Module();
					mod.setName((new String(buffer)).trim());
					modules.add(mod);
				}
			}
		} else {
			return new ArrayList<Module>(); // TODO: should maybe throw an exception if modules were not able to be retrieved
		}
		
		kernel32.CloseHandle(hProcess);
		
		return modules;
	}
	
	
	
	public static List<SystemProcess> fetchProcesses() throws Exception {
		List<SystemProcess> processes = new ArrayList<SystemProcess>();
		
		Kernel32 kernel32 = Kernel32.INSTANCE;
		
		// retrieve all processes
		HANDLE snapshot = kernel32.CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
		int handleValue = (new PointerByReference(snapshot.getPointer())).getPointer().getInt(0);
		
		if (handleValue == -1) { // INVALID_HANDLE_VALUE
			throw new Exception("Could not retrieve process list: invalid handle returned.");
		}
		else {
			LPPROCESSENTRY32 procEntry = new LPPROCESSENTRY32();
			boolean success = kernel32.Process32First(snapshot, procEntry);
			if(!success) {
				throw new Exception("Process32First: Could not get first process from list. Error:" + kernel32.GetLastError());
			}
			
			// iterate over list
			do {
				SystemProcess proc = new SystemProcess();
				proc.setName(procEntry.getSzExeFile());
				proc.setId(procEntry.th32ProcessID);
				proc.setParentId(procEntry.th32ParentProcessID);
				
				processes.add(proc);
			} while (kernel32.Process32Next(snapshot, procEntry));
		}
		
		return processes;
	}
	
	
}
